{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.99833065,  0.05775759, -0.778794  ], dtype=float32)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "    def __init__(self):\n",
    "        env = gym.make('Pendulum-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        done = terminated or truncated\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            done = True\n",
    "        return state, reward, done, info\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200,\n",
       " ([0.6064112186431885, -0.7951512336730957, -0.83575838804245],\n",
       "  0.5332061648368835,\n",
       "  -0.9151667857748363,\n",
       "  [0.5513089895248413, -0.8343011736869812, -1.3521409034729004],\n",
       "  False),\n",
       " torch.Size([200, 4]),\n",
       " torch.Size([200, 4]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "\n",
    "class Pool:\n",
    "    def __init__(self, limit):\n",
    "        #样本池\n",
    "        self.datas = []\n",
    "        self.limit = limit\n",
    "\n",
    "    def add(self, state, action, reward, next_state, over):\n",
    "        if isinstance(state, np.ndarray) or isinstance(state, torch.Tensor):\n",
    "            state = state.reshape(3).tolist()\n",
    "\n",
    "        action = float(action)\n",
    "\n",
    "        reward = float(reward)\n",
    "\n",
    "        if isinstance(next_state, np.ndarray) or isinstance(\n",
    "                next_state, torch.Tensor):\n",
    "            next_state = next_state.reshape(3).tolist()\n",
    "\n",
    "        over = bool(over)\n",
    "\n",
    "        self.datas.append((state, action, reward, next_state, over))\n",
    "        #数据上限,超出时从最古老的开始删除\n",
    "        while len(self.datas) > self.limit:\n",
    "            self.datas.pop(0)\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def get_sample(self):\n",
    "        #从样本池中采样\n",
    "        samples = self.datas\n",
    "\n",
    "        #[b, 3]\n",
    "        state = torch.FloatTensor([i[0] for i in samples]).reshape(-1, 3).to(device='cuda')\n",
    "        #[b, 1]\n",
    "        action = torch.FloatTensor([i[1] for i in samples]).reshape(-1, 1).to(device='cuda')\n",
    "        #[b, 1]\n",
    "        reward = torch.FloatTensor([i[2] for i in samples]).reshape(-1, 1).to(device='cuda')\n",
    "        #[b, 3]\n",
    "        next_state = torch.FloatTensor([i[3] for i in samples]).reshape(-1, 3).to(device='cuda')\n",
    "        #[b, 1]\n",
    "        over = torch.LongTensor([i[4] for i in samples]).reshape(-1, 1).to(device='cuda')\n",
    "\n",
    "        #[b, 4]\n",
    "        input = torch.cat([state, action], dim=1)\n",
    "        #[b, 4]\n",
    "        label = torch.cat([reward, next_state - state], dim=1)\n",
    "\n",
    "        return input, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.datas)\n",
    "\n",
    "\n",
    "pool = Pool(100000)\n",
    "\n",
    "\n",
    "#初始化一局游戏的数据\n",
    "def _():\n",
    "    #初始化游戏\n",
    "    state = env.reset()\n",
    "\n",
    "    #玩到游戏结束为止\n",
    "    over = False\n",
    "    while not over:\n",
    "        #随机一个动作\n",
    "        action = env.action_space.sample()[0]\n",
    "\n",
    "        #执行动作,得到反馈\n",
    "        next_state, reward, over, _ = env.step([action])\n",
    "\n",
    "        #记录数据样本\n",
    "        pool.add(state, action, reward, next_state, over)\n",
    "\n",
    "        #更新游戏状态,开始下一个动作\n",
    "        state = next_state\n",
    "\n",
    "\n",
    "_()\n",
    "\n",
    "a, b = pool.get_sample()\n",
    "\n",
    "len(pool), pool.datas[0], a.shape, b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Gym",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
